import json
import time
from train import StartTrain
import pymysql

def trainONNX():
    start = time.time()
    DB,_ = connectDB()
    StartTrain(DB)

    # 关闭连接对象，否则会导致连接泄漏，消耗数据库资源
    DB.close()
    # 关闭光标
    DB.close()

    trainTime = time.time() - start
    Output = "训练完成，耗时" + str("{:.2f}".format(trainTime)) + "S"
    print(Output)
    return Output

def connectDB():

    with open("./data/DBconfig.json", "r") as f:
        config = json.load(f)
    host = config["host"]
    port = config["port"]
    user = config["user"]
    password = config["password"]
    db = config["db"]
    charset = config["charset"]

    # 连接数据库
    connection = pymysql.connect(host=host, # host属性
    port=port,  # 端口，默认为3306
    user = user, # 用户名
    password = password, # 此处填登录数据库的密码
    db = db, # 数据库名
    charset = charset,  # 字符编码
    )
    cur = connection.cursor()

    return cur,connection

trainONNX()

